import torch
import config
import torch.nn as nn


class YoloV1Loss (nn.Module) :
    def __init__ (self, batch_size=config.BATCH_SIZE, S=config.S, B=config.B, C=config.C, l_coord=5, l_noobj=0.5) :
        super (YoloV1Loss, self).__init__ ()
        self.S = S
        self.B = B
        self.C = C
        self.depth = self.B * 5 + self.C
        self.batch_size = batch_size
        self.l_coord = l_coord
        self.l_noobj = l_noobj

    def forward (self, preds, labels) :
        """
        :brief calculate the loss between preds and labels
        :param preds: (batch, S, S, depth) by default (128, 7, 7, 30)
        :param labels: (batch, S, S, depth) by default (128, 7, 7, 30)
        :return: return the Sum Square Loss of Yolo
        """
        coord_loss = self.cal_coord_loos (preds, labels)
        bbox_loss = self.cal_bbox_loss (preds, labels)
        obj_conf_loss = self.cal_obj_conf_loss (preds, labels)
        noobj_conf_loss = self.cal_noobj_conf_loss (preds, labels)
        class_loss = self.cal_class_loss (preds, labels)
        total_loss = self.l_coord * coord_loss + self.l_coord * bbox_loss + \
                     obj_conf_loss + self.l_noobj * noobj_conf_loss + class_loss
        loss_of_each_data = torch.sum (total_loss) / self.batch_size
        return loss_of_each_data

    def cal_coord_loos (self, pred, label) :
        """
        calculate the coordinate square error
        :param pred: the predictions
        :param label: the labels
        :return: the loss with bacth information, the shape is (batch_size, 1)
        """
        loss = torch.zeros ((self.batch_size, 1)).to ('cuda')

        ## calculate by for loop
        # loss.shape : (128, 1)
        # for i in range (self.S) :
        #     for j in range (self.S) :
        #         if label[:, i, j, 0] == 1 :
        #             loss += torch.sum (pred[:, i, j, 1] - label[:, i, j, 1]) ** 2 \
        #             + (pred[:, i, j, 2] - label[:, i, j, 2]) ** 2)

        # calculate by vector operation
        loss += torch.sum (
            label[:, :, :, 0] *
            (
                    (label[:, :, :, 1] - pred[:, :, :, 1]) ** 2 +
                    (label[:, :, :, 2] - pred[:, :, :, 2]) ** 2
            )
        ) + torch.sum (
            label[:, :, :, 5] *
            (
                    (label[:, :, :, 6] - pred[:, :, :, 6]) ** 2 +
                    (label[:, :, :, 7] - pred[:, :, :, 7]) ** 2
            )
        )

        # label shape : (batch_size, S, S, depth); label[:, :, :, 0] shape : (batch_size, S, S)
        # loss shape : (batch_size, 1)

        return loss

    def cal_bbox_loss (self, pred, label) :
        """
        calculate the bbox square error
        :param pred: the predictions
        :param label: the labels
        :return: the loss with batch information, the shape is (batch_size, 1)
        """
        loss = torch.zeros ((self.batch_size, 1)).to ('cuda')

        loss += torch.sum (  # first bbox in all grid cells
            label[:, :, :, 0] *
            (
                    (label[:, :, :, 3] - pred[:, :, :, 3]) ** 2
                    + (label[:, :, :, 4] - pred[:, :, :, 4]) ** 2
            )
        ) + torch.sum (  # second bbox in all grid cells
            label[:, :, :, 5] *
            (
                    (label[:, :, :, 8] - pred[:, :, :, 8]) ** 2
                    + (label[:, :, :, 9] - pred[:, :, :, 9]) ** 2
            )
        )

        return loss

    def cal_obj_conf_loss (self, pred, label) :
        """
        calculate the obj confidence loss
        :param pred: the predictions
        :param label: the labels
        :return: the loss with batch information, the shape is (batch_size, 1)
        """
        loss = torch.zeros ((self.batch_size, 1)).to ('cuda')

        loss += torch.sum (
            label[:, :, :, 0] * (label[:, :, :, 0] - pred[:, :, :, 0]) ** 2
        )

        return loss

    def cal_noobj_conf_loss (self, pred, label) :
        """
        calculate the noobj confidence loss
        :param pred: the predictions
        :param label: the labels
        :return: the loss with batch information, the shape is (batch_size, 1)
        """
        loss = torch.zeros ((self.batch_size, 1)).to ('cuda')

        loss += torch.sum (
            (-(label[:, :, :, 0] - 1)) * (label[:, :, :, 0] - pred[:, :, :, 0]) ** 2
        )
        # -(label[:, :, : 0] - 1) will reverse the 0 and 1

        return loss

    def cal_class_loss (self, pred, label) :
        """
        calculate the prediction class error
        :param pred: the predictions
        :param label: the labels
        :return: the loss with batch information, the shape is (batch_size, 1)
        """
        loss = torch.zeros ((self.batch_size, 1)).to ('cuda')

        loss += torch.sum (
            label[:, :, :, 0] * (
                torch.sum ((label[:, :, :, 10 :30] - pred[:, :, :, 10 :30]) ** 2, dim=3)
            )
        )
        # go through the 20 classes when the obj exits
        return loss


if __name__ == '__main__' :
    pred = torch.zeros ((2, 3, 3, 5))
    label = torch.randn ((2, 3, 3, 5))
    loss = torch.zeros ((2, 1))
    print (loss.shape)
    print (torch.sum (pred).shape)
    loss += torch.sum (pred[:, :, :, 0] - label[:, :, :, 0])
    print (loss)
